{-# OPTIONS -XTypeOperators #-}

module FTable (
   (:->:)(..),
   liftFT, liftFT2, liftFT3,
   fromList, fromList2, fromList3,
   toList, toList2, toList3,
   assocs, assocs2, assocs3,
   elems, elems2, elems3,
   mapWithArg,
   quad, interpolate, extrapolate,
   translate, translate2,
   move, move2,
   scale, scale2,

   ap, unsafeInterpolate

) where

-- import Debug.Trace

import Control.Applicative
import Control.Monad ( replicateM )
import Data.Binary
import Data.Map (Map)
import Data.Ix

import qualified Data.Array as A
import qualified Data.Map   as M

import NumInstances

data i :->: a = F (i,i,i) (Map i a) | C a

---- Экспортируемые функции

liftFT :: (Integral i, Num a) => (i,i,i) -> (i->a) -> i :->: a
liftFT g@(a,b,d) f
   | a > b              = error "В функции liftFT должно выполняться a <= b"
   | d <= 0             = error "В функции liftFT должно выполняться d > 0"
   | (b-a) `mod` d /= 0 = error "В функции liftFT должно выполняться (b-a) `mod` d == 0"
   | otherwise          = F g $ M.fromList [(x, f x) | x <- [a, a+d .. b]]

liftFT2 (g1,g2)    f = liftFT g1 $ \x -> liftFT g2  (f x)
liftFT3 (g1,g2,g3) f = liftFT g1 $ \x -> liftFT g2 $ \y -> liftFT g3 (f x y)

fromList g xs = liftFT g f
   where
      f x = (M.fromList xs) M.! x

fromList2 g xs = liftFT2 g f
   where
      f x y = (M.fromList xs) M.! (x,y)

fromList3 g xs = liftFT3 g f
   where
      f x y z = (M.fromList xs) M.! (x,y,z)


toList :: Integral i => i :->: a -> [(i,a)]
toList (F _ m) = M.toList m

toList2 f = [((x,y), u) |
               (x, f') <- toList f ,
               (y, u ) <- toList f']

toList3 f = [((x,y,z), u) |
               ( x,    f') <- toList  f ,
               ((y,z), u ) <- toList2 f']

assocs  = toList
assocs2 = toList2
assocs3 = toList3

elems :: Integral i => i :->: a -> [a]
elems (F _ m) = M.elems m

elems2 f = concatMap elems $ elems f
elems3 f = concatMap elems $ elems2 f

mapWithArg :: (i -> a -> b) -> (i :->: a) -> (i :->: b)
mapWithArg f (F g m) = F g $ M.mapWithKey f m

translate :: Integral i => i -> (i :->: a) -> (i :->: a)
translate v (F (a,b,d) m) = let v' = fromIntegral v in F (a+v',b+v',d) $ M.mapKeys (v+) m
translate _ f = f

move :: (RealFrac r, Integral i, Fractional a) => r -> (i :->: a) -> (i :->: a)
move v f@(F g@(_,_,d) _)
      | abs v <= 1e-15 = f
      | abs v <  d'    = liftFT g (\x -> extrapolate f (fromIntegral x + v))
      | otherwise      = move (v - fromIntegral v') $ translate v' f
      where
         d' = fromIntegral d
         v' = d * floor (v/d')
move _ f = f

scale :: (Integral i, Fractional a, RealFrac r) => r -> (i :->: a) -> (i :->: a)
scale s f@(F (a,b,d) m)
   | s == 1 = f
   | s <  1 = F (floor a', floor b', ceiling d') m'
   | otherwise = error "В функции scale должно выполняться s < 1"
   where
      a' = s * fromIntegral a
      b' = s * fromIntegral b
      d' = s * fromIntegral d 
      m' = M.fromListWith (+) $ concatMap prepare $ M.assocs m
      prepare (x,v)
         | lb == rb  = [(lb, v)]
         | otherwise = [(lb, v * k), (rb, v * (1-k))]
         where
            k  = fromRational $ toRational $ (fromIntegral rb - y) / d'
            lb = floor (y)
            rb = floor (y+d')
            y  = s * fromIntegral x

scale s f = f / (fromRational $ toRational s)

translate2   (x,y)   = (translate x) . (fmap $ translate y)
move2        (x,y)   = (move      x) . (fmap $ move      y)
scale2       (x,y)   = (scale     x) . (fmap $ scale     y)

translate3   (x,y,z) = (translate x) . (fmap $ translate2 (y,z))
move3        (x,y,z) = (move      x) . (fmap $ move2      (y,z))
scale3       (x,y,z) = (scale     x) . (fmap $ scale2     (y,z))

quad ::  (Integral i, Fractional a) => (i :->: a) -> a
quad f@(F (a,b,d) m) = (fromIntegral d) * (e - c)
   where
      c = 0.5 * (f `ap` a + f `ap` b)
      e = sum $ M.elems m

extrapolate, interpolate :: (Integral i, RealFrac a, Fractional b) => (i :->: b) -> a -> b

extrapolate f@(F (a,b,_) _) x
   | x <= a'   = f `ap` a
   | x >= b'   = f `ap` b
   | otherwise = unsafeInterpolate f x
   where
      a' = fromIntegral a
      b' = fromIntegral b

extrapolate f x = unsafeInterpolate f x

interpolate f@(F (a,b,_) _) x
   | x < a'    = 0
   | x > b'    = 0
   | otherwise = unsafeInterpolate f x
   where
      a' = fromIntegral a
      b' = fromIntegral b

interpolate f x = unsafeInterpolate f x


---- Наследование типа данных (:->:) от нужных классов

instance (Show i, Show b) => Show (i :->: b) where
   show (F g  _) = show g
   show (C v)    = show v

instance (Eq i, Eq b) => Eq (i :->: b) where
   (==) =  noFun "(==)"
   (/=) =  noFun "(/=)"

instance Integral i => Functor ( (:->:) i ) where
   fmap f x = pure f <*> x

instance Integral i => Applicative ( (:->:) i ) where
   pure a = C a

   (<*>) (C f)   (C a)   = C (f a)
   (<*>) (C f)   (F g m) = F g (M.map f m)
   (<*>) (F g m) (C a)   = F g (M.map (flip id a) m)
   (<*>) (F _ _) (F _ _) = undefined

instance (Integral i, Num a) => Num (i :->: a) where
   negate      = fmap negate
   (*)         = (*.)
   (+)         = (+.)
   fromInteger = pure . fromInteger
   abs         = fmap abs
   signum      = fmap signum

instance (Integral i, Fractional a) => Fractional (i :->: a) where
   recip        = fmap recip
   fromRational = pure . fromRational

instance (Binary i, Integral i, Binary b) => Binary (i :->: b) where
   put (C v)     =  do put (1 :: Word8)
                       put v
   put (F g m)   =  do put (0 :: Word8)
                       put g
                       put n       -- write the length
                       mapM_ put e -- now the elems
                       where
                          e = M.assocs m
                          n = length e

   get = do t <- get :: Get Word8
            case t of
               1 -> do v <- get
                       return (C v)
               0 -> do g <- get
                       n <- get                -- read the length
                       xs <- replicateM n get  -- now the elems.
                       let m = M.fromList xs
                       return (F g m)

---- Не экспортируемые функции

noOv :: String -> String -> a
noOv ty meth = error $ meth ++ ": No overloading for " ++ ty

noFun :: String -> a
noFun = noOv "function"

ap :: (Integral i, Num a) => (i :->: a) -> i -> a
ap (C a) _ = a
ap (F (a,b,d) m) x | x < a || x > b       = 0
                   | (x - a) `mod` d /= 0 = error "В функции ap должно (x - a) `mod` d /= 0"
                   | otherwise            = m M.! x

validateGrid :: Integral a => (a,a,a) -> (a,a,a) -> Bool
-- validateGrid a b | trace ("validateGrid " ++ show a ++ " " ++ show b) False = undefined
validateGrid (a,_,d) (a',_,d') = d /= d' || a' /= a + d * ( (a'-a) `div` d)

intersection, union :: Integral a => (a,a,a) -> (a,a,a) -> (a,a,a)
intersection (a,b,_) (c,d,e) = ( max a c, min b d, e )
union        (a,b,_) (c,d,e) = ( min a c, max b d, e )

(*.), (+.) :: (Integral i, Num a) => (i :->: a) -> (i :->: a) -> (i :->: a)

(*.) (F g m) (F g' m')
   | validateGrid g g' = error "Функции заданны на несоответствующих друг другу сетках"
   | M.null m''                 = C 0 -- TODO : Проверить, а это вообще правильно так делать?
   | otherwise                  = F g'' m''
   where
      m'' = M.intersectionWith (*) m m'
      g'' = intersection g g'

(*.) a b = liftA2 (*) a b

(+.) (F g m) (F g' m')
   | validateGrid g g' = error "Функции заданны на несоответствующих друг другу сетках"
   | otherwise                  = F g'' m''
   where
      m'' = M.unionWith (+) m m'
      g'' = union g g'

(+.) a b = liftA2 (+) a b


unsafeInterpolate :: (Integral i, RealFrac a, Fractional b) => (i :->: b) -> a -> b
unsafeInterpolate f@(F (a,b,d) _) x = v0 + k * (v1 - v0)
   where
      v0 = f `ap` x0
      v1 = f `ap` x1
      x0 = a + d * ( ceiling $ ( x - a' ) / d' - 1 )
      x1 = x0 + d
      k  = fromRational $ toRational $ (x - fromIntegral x0) / d'

      a' = fromIntegral a
      d' = fromIntegral d

unsafeInterpolate (C v) _ = v